from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import Optional, Dict, Any

from rag_service import RAGService


class QueryRequest(BaseModel):
    question: str


class QueryResponse(BaseModel):
    answer: str
    sources: list[dict]


class RetrievalRequest(BaseModel):
    query: str
    parameters: Optional[Dict[str, Any]] = None


class RetrievalResponse(BaseModel):
    results: list[dict]


app = FastAPI(
    title="DITING RAG API",
    version="0.1.0",
    description="Simple RAG endpoints exposing retrieval-augmented answers.",
)

rag_service = RAGService()


@app.get("/health")
def health_check() -> dict:
    return {"status": "ok"}


@app.post("/rag/query", response_model=QueryResponse)
async def rag_query(payload: QueryRequest) -> QueryResponse:
    try:
        answer, sources = rag_service.query(payload.question)
    except ValueError as err:
        raise HTTPException(status_code=400, detail=str(err)) from err
    except Exception as err:  # pragma: no cover
        raise HTTPException(status_code=500, detail="RAG pipeline failure") from err

    return QueryResponse(answer=answer, sources=sources)


@app.post("/api/v1/rag/retrieval", response_model=RetrievalResponse)
async def rag_retrieval(payload: RetrievalRequest) -> RetrievalResponse:
    # 获取参数
    parameters = payload.parameters or {}
    similarity_threshold = parameters.get("similarity_threshold", 0.0)
    context_recall_max_tokens = parameters.get("context_recall_max_tokens", 4000)
    k = parameters.get("top_k", rag_service.settings.search_k)
    
    try:
        # 执行检索
        results = rag_service.vectorstore.similarity_search_with_score(
            payload.query, k=k
        )
        
        # 根据相似度阈值过滤结果
        filtered_results = []
        total_tokens = 0
        
        for doc, score in results:
            # 相似度阈值过滤
            if score < similarity_threshold:
                continue
                
            # Token数量限制
            doc_tokens = len(doc.page_content.split())
            if total_tokens + doc_tokens > context_recall_max_tokens:
                break
                
            filtered_results.append((doc, score))
            total_tokens += doc_tokens
        
        # 格式化结果
        formatted_results = [
            {
                "content": doc.page_content,
                "metadata": dict(doc.metadata),
                "similarity": float(score)
            }
            for doc, score in filtered_results
        ]
        
    except Exception as err:
        raise HTTPException(status_code=500, detail=f"Retrieval failed: {str(err)}") from err

    return RetrievalResponse(results=formatted_results)


if __name__ == "__main__":
    import uvicorn
    uvicorn.run("api:app", host="0.0.0.0", port=8000, reload=True)